package com.pai4j.message.websocket;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketSession;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket会话管理器
 * 管理不同业务场景下的WebSocket Session
 * 
 * 支持多端推送（PC、Mobile等）
 * 线程安全的Session管理
 *
 * @author PAI4J Team
 * @date 2025-11-01
 */
@Slf4j
@Component
public class WebSocketSessionManager {

    /**
     * Session存储结构
     * Map<bizCode, Map<userId, WebSocketSession>>
     * 
     * bizCode：业务代码，用于区分不同场景（如：PC_CHAT、MOBILE_CHAT）
     * userId：用户ID
     * WebSocketSession：WebSocket会话对象
     */
    private static final Map<String, Map<String, WebSocketSession>> sessions = new ConcurrentHashMap<>();

    /**
     * 注册Session
     *
     * @param bizCode 业务代码
     * @param userId 用户ID
     * @param session WebSocket会话
     */
    public void registerSession(String bizCode, String userId, WebSocketSession session) {
        if (StringUtils.isAnyBlank(bizCode, userId) || session == null) {
            log.warn("注册Session失败：参数不能为空");
            return;
        }

        Map<String, WebSocketSession> userSessions = sessions.computeIfAbsent(bizCode, k -> new ConcurrentHashMap<>());

        // 如果同一bizCode下该用户已有连接，关闭旧连接，防止重复连接堆积
        WebSocketSession oldSession = userSessions.get(userId);
        if (oldSession != null && oldSession.isOpen() && !StringUtils.equals(oldSession.getId(), session.getId())) {
            try {
                log.warn("检测到重复连接，关闭旧会话。userId:{}, bizCode:{}, oldSessionId:{}, newSessionId:{}", 
                        userId, bizCode, oldSession.getId(), session.getId());
                oldSession.close();
            } catch (Exception e) {
                log.warn("关闭旧会话失败。userId:{}, oldSessionId:{}", userId, oldSession.getId(), e);
            }
        }

        userSessions.put(userId, session);
        log.info("WebSocket会话已注册 => bizCode:{}, userId:{}, sessionId:{}", bizCode, userId, session.getId());
    }

    /**
     * 获取单个用户的Session
     *
     * @param bizCode 业务代码
     * @param userId 用户ID
     * @return WebSocket会话
     */
    public WebSocketSession getSession(String bizCode, String userId) {
        if (StringUtils.isAnyBlank(bizCode, userId)) {
            return null;
        }

        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (MapUtils.isEmpty(userSessions)) {
            return null;
        }

        return userSessions.get(userId);
    }

    /**
     * 批量获取Session（多对多群聊场景）
     *
     * @param bizCode 业务代码
     * @param userIds 用户ID集合
     * @return 用户ID -> WebSocket会话的映射
     */
    public Map<String, WebSocketSession> getSessions(String bizCode, Set<String> userIds) {
        if (StringUtils.isBlank(bizCode) || userIds == null || userIds.isEmpty()) {
            return Collections.emptyMap();
        }

        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (MapUtils.isEmpty(userSessions)) {
            return Collections.emptyMap();
        }

        Map<String, WebSocketSession> result = new ConcurrentHashMap<>();
        userIds.forEach(userId -> {
            WebSocketSession session = userSessions.get(userId);
            if (session != null && session.isOpen()) {
                result.put(userId, session);
            }
        });

        return result;
    }

    /**
     * 获取用户在所有业务场景下的Session
     *
     * @param userId 用户ID
     * @return 业务代码 -> WebSocket会话的映射
     */
    public Map<String, WebSocketSession> getAllSessionsByUser(String userId) {
        if (StringUtils.isBlank(userId)) {
            return Collections.emptyMap();
        }

        Map<String, WebSocketSession> result = new ConcurrentHashMap<>();
        sessions.forEach((bizCode, userSessions) -> {
            WebSocketSession session = userSessions.get(userId);
            if (session != null && session.isOpen()) {
                result.put(bizCode, session);
            }
        });

        return result;
    }

    /**
     * 移除Session
     *
     * @param bizCode 业务代码
     * @param userId 用户ID
     */
    public void removeSession(String bizCode, String userId) {
        if (StringUtils.isAnyBlank(bizCode, userId)) {
            return;
        }

        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (MapUtils.isEmpty(userSessions)) {
            return;
        }

        WebSocketSession removed = userSessions.remove(userId);
        if (removed != null) {
            log.info("WebSocket会话已移除 => bizCode:{}, userId:{}, sessionId:{}", bizCode, userId, removed.getId());
        }

        // 如果该bizCode下已无会话，移除整个Map
        if (userSessions.isEmpty()) {
            sessions.remove(bizCode);
        }
    }

    /**
     * 移除用户在特定Session下的连接
     *
     * @param bizCode 业务代码
     * @param userId 用户ID
     * @param sessionId 会话ID
     */
    public void removeSessionById(String bizCode, String userId, String sessionId) {
        if (StringUtils.isAnyBlank(bizCode, userId, sessionId)) {
            return;
        }

        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (MapUtils.isEmpty(userSessions)) {
            return;
        }

        WebSocketSession existing = userSessions.get(userId);
        if (existing != null && StringUtils.equals(existing.getId(), sessionId)) {
            userSessions.remove(userId);
            log.info("WebSocket会话已移除（按SessionId） => bizCode:{}, userId:{}, sessionId:{}", bizCode, userId, sessionId);
            
            if (userSessions.isEmpty()) {
                sessions.remove(bizCode);
            }
        }
    }

    /**
     * 检查用户是否在线（任意业务场景下）
     *
     * @param userId 用户ID
     * @return true-在线，false-离线
     */
    public boolean isOnline(String userId) {
        if (StringUtils.isBlank(userId)) {
            return false;
        }

        return sessions.values().stream()
                .anyMatch(userSessions -> {
                    WebSocketSession session = userSessions.get(userId);
                    return session != null && session.isOpen();
                });
    }

    /**
     * 检查用户在特定业务场景下是否在线
     *
     * @param bizCode 业务代码
     * @param userId 用户ID
     * @return true-在线，false-离线
     */
    public boolean isOnline(String bizCode, String userId) {
        if (StringUtils.isAnyBlank(bizCode, userId)) {
            return false;
        }

        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (MapUtils.isEmpty(userSessions)) {
            return false;
        }

        WebSocketSession session = userSessions.get(userId);
        return session != null && session.isOpen();
    }

    /**
     * 获取所有业务代码
     *
     * @return 业务代码集合
     */
    public Set<String> getAllBizCodes() {
        return new HashSet<>(sessions.keySet());
    }

    /**
     * 获取指定业务代码下的在线用户数
     *
     * @param bizCode 业务代码
     * @return 在线用户数
     */
    public int getOnlineUserCount(String bizCode) {
        if (StringUtils.isBlank(bizCode)) {
            return 0;
        }

        Map<String, WebSocketSession> userSessions = sessions.get(bizCode);
        if (MapUtils.isEmpty(userSessions)) {
            return 0;
        }

        return (int) userSessions.values().stream()
                .filter(session -> session != null && session.isOpen())
                .count();
    }

    /**
     * 获取总在线用户数（所有业务场景）
     *
     * @return 总在线用户数
     */
    public int getTotalOnlineUserCount() {
        Set<String> uniqueUsers = new HashSet<>();
        sessions.values().forEach(userSessions -> 
            userSessions.forEach((userId, session) -> {
                if (session != null && session.isOpen()) {
                    uniqueUsers.add(userId);
                }
            })
        );
        return uniqueUsers.size();
    }

    /**
     * 清理所有失效的Session
     *
     * @return 清理的Session数量
     */
    public int cleanupInvalidSessions() {
        int cleaned = 0;
        
        Iterator<Map.Entry<String, Map<String, WebSocketSession>>> bizIterator = sessions.entrySet().iterator();
        while (bizIterator.hasNext()) {
            Map.Entry<String, Map<String, WebSocketSession>> bizEntry = bizIterator.next();
            String bizCode = bizEntry.getKey();
            Map<String, WebSocketSession> userSessions = bizEntry.getValue();

            Iterator<Map.Entry<String, WebSocketSession>> userIterator = userSessions.entrySet().iterator();
            while (userIterator.hasNext()) {
                Map.Entry<String, WebSocketSession> userEntry = userIterator.next();
                WebSocketSession session = userEntry.getValue();
                
                if (session == null || !session.isOpen()) {
                    userIterator.remove();
                    cleaned++;
                    log.debug("清理失效会话 => bizCode:{}, userId:{}", bizCode, userEntry.getKey());
                }
            }

            if (userSessions.isEmpty()) {
                bizIterator.remove();
            }
        }

        if (cleaned > 0) {
            log.info("清理失效WebSocket会话完成，共清理{}个", cleaned);
        }

        return cleaned;
    }

    /**
     * 获取所有Session统计信息
     *
     * @return 统计信息
     */
    public Map<String, Object> getStatistics() {
        Map<String, Object> stats = new HashMap<>();
        stats.put("totalBizCodes", sessions.size());
        stats.put("totalOnlineUsers", getTotalOnlineUserCount());
        
        Map<String, Integer> bizCodeStats = new HashMap<>();
        sessions.forEach((bizCode, userSessions) -> 
            bizCodeStats.put(bizCode, getOnlineUserCount(bizCode))
        );
        stats.put("bizCodeStats", bizCodeStats);
        
        return stats;
    }
}

